from sklearn.model_selection import train_test_split  # 导入切分训练集、测试集模块
from sklearn.linear_model import LinearRegression
import getData_ML
from numpy import array
import time

import get_closeMongoDB


def predict_data():
    x_data, y_data, precict_data = getData_ML.getAddHistory()
    X = array(x_data)
    y = array(y_data)
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.06, random_state=532)

    print(X_test)
    linreg = LinearRegression()
    # 训练
    model = linreg.fit(X_train, y_train)
    # print('模型参数:')
    # print(model)
    # # 训练后模型截距
    # print('模型截距:')
    # print(linreg.intercept_)
    # # 训练后模型权重（特征个数无变化）
    # print('回归系数:')
    # print(linreg.coef_)

    # y_pred = linreg.predict(X_test)
    # print("拟合优度：")
    print("拟合优度：" + str(linreg.score(X_train, y_train)))
    predict_last = linreg.predict(array(precict_data))
    # print(int(predict_last))
    return int(predict_last)


def get_predict_data():
    # res = predict_data()
    Res_time = time.strftime("%Y-%m-%d", time.localtime())
    q_rs = get_closeMongoDB.getRegisContent(Res_time)  # 获取查询结果
    if q_rs is not None:  # 不为空则返回（一天只预测一次）
        del q_rs['_id']
        return q_rs
    else:
        insert_content = predict_data()  # 获取预测数据
        insertRs = {"time": Res_time, "data": insert_content}
        get_closeMongoDB.insertRegisContent(insertRs)  # 插入预测数据
        print("insertRegisContent ")
        del insertRs['_id']
        return insertRs
